package com.simple.llm.service;

import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.github.zuihou.base.R;
import com.simple.common.exception.GlobalException;
import com.simple.llm.config.LLMProperties;
import com.simple.llm.domain.dto.LLMMessageDTO;
import com.simple.llm.domain.dto.LLMStreamResultDTO;
import com.simple.llm.domain.dto.NormalChatDTO;
import com.simple.llm.domain.dto.StreamChatDTO;
import com.simple.llm.domain.entity.LLMRecords;
import com.simple.llm.domain.entity.LLMUsageRecords;
import com.simple.llm.domain.vo.LLMResponseVO;
import com.simple.llm.mapper.LLmRecordsMapper;
import com.simple.llm.mapper.LLmUsageRecordsMapper;
import com.simple.llm.service.impl.DouBaoLLMStrategyImpl;
import com.simple.llm.service.impl.GLMLLMStrategyImpl;
import com.simple.llm.service.impl.OpenAiInterfaceLLMStrategyImpl;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletResponse;
import java.util.*;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@Slf4j
@Component
@RefreshScope
public class LLMContext implements LLMStrategy, InitializingBean{
    @Resource
    private LLMProperties llmProperties;
    @Resource
    private LLmUsageRecordsMapper lLmUsageRecordsMapper;
    @Resource
    private LLmRecordsMapper lLmRecordsMapper;

    @Resource
    private ObjectMapper objectMapper;
    private final Map<String, LLMStrategy> strategyMap = new HashMap<>();

    private String defaultModelName = "";

    private void initStrategy() {
        List<LLMProperties.GptProperties> gptList = llmProperties.getGptList();
        strategyMap.clear();
        for (LLMProperties.GptProperties gptProperties : gptList) {
            LLMStrategy impl = null;
            if (StringUtils.isBlank(gptProperties.getInterfaceName()) || gptProperties.getInterfaceName().equalsIgnoreCase("openai")) {
                impl = new OpenAiInterfaceLLMStrategyImpl(gptProperties,objectMapper);
            } else if (gptProperties.getInterfaceName().equalsIgnoreCase("glm")) {
                impl = new GLMLLMStrategyImpl(gptProperties);
            }else if (gptProperties.getInterfaceName().equalsIgnoreCase("doubao")) {
                impl = new DouBaoLLMStrategyImpl(gptProperties);
            }
            // 预留的其它实现初始化

            if (impl == null) {
                throw new Error("存在未实现接口的大模型配置:" + gptProperties);
            }
            strategyMap.put(gptProperties.getModelName(), impl);
            if (Boolean.TRUE.equals(gptProperties.getDefaultLLM())) {
                defaultModelName = gptProperties.getModelName();
                strategyMap.put("default", impl);
            }
        }
        if (strategyMap.size() == 0) {
            throw new Error("当前未配置任务llm，服务启动失败");
        }
    }

    private LLMStrategy getLLM(NormalChatDTO normalChatDTO) {
        String modelName = normalChatDTO.getModelName();
        if (StringUtils.isBlank(modelName)) {
            modelName = "default";
        }
        LLMStrategy llmStrategy = strategyMap.get(modelName);
        if (llmStrategy == null) {
            throw new GlobalException("未找到模型:" + modelName);
        }
        // 如果存在排序，则对messages进行排序
        List<LLMMessageDTO> collect = normalChatDTO.getMessages().stream().sorted((Comparator.comparing(LLMMessageDTO::getIndex))).collect(Collectors.toList());
        normalChatDTO.setMessages(collect);
        return llmStrategy;
    }

    public List<LLMMessageDTO> getMessagesByChatId(String chatId) {
        LLMRecords llmRecords = lLmRecordsMapper.selectOne(Wrappers.<LLMRecords>lambdaQuery()
                .eq(LLMRecords::getChatId, chatId));
        if(llmRecords == null){
            return new LinkedList<>();
        }
        String inputContent = llmRecords.getInputContent();
        List<LLMMessageDTO> llmMessageDTOS = JSON.parseArray(inputContent, LLMMessageDTO.class);
        if(StringUtils.isNotBlank(llmRecords.getOutputContent())){
            llmMessageDTOS.add(LLMMessageDTO.assistantMsg(llmRecords.getOutputContent(),llmMessageDTOS.size()));
        }
        return llmMessageDTOS.stream().sorted((Comparator.comparing(LLMMessageDTO::getIndex))).collect(Collectors.toList());
    }

    public String streamRecord(StreamChatDTO streamChatDTO){
        String chatId = record(streamChatDTO, null, null, true);
        streamChatDTO.setChatId(chatId);
        return chatId;
    }
    public String record(NormalChatDTO normalChatDTO, LLMResponseVO vo, String chatId, boolean stream) {
        if (chatId == null) {
            // 如果消息记录没有序号就给一个，保证序列化时不丢失排序
            for (int i = 0; i < normalChatDTO.getMessages().size(); i++) {
                LLMMessageDTO llmMessageDTO = normalChatDTO.getMessages().get(i);
                if(llmMessageDTO.getIndex() == null){
                    llmMessageDTO.setIndex(i);
                }
            }

            chatId = UUID.randomUUID().toString().replace("-", "");
            LLMRecords build = LLMRecords.builder()
                    .chatId(chatId)
                    .inputContent(JSON.toJSONString(normalChatDTO.getMessages()))
                    .updateTime(new Date())
                    .createTime(normalChatDTO.getCreateDate())
                    .createUser(normalChatDTO.getCreateUser())
                    .streamChat(stream)
                    .build();
            lLmRecordsMapper.insert(build);
        }
        if (vo != null) {
            LLMUsageRecords records = LLMUsageRecords.builder()
                    .chatId(chatId)
                    .createUser(vo.getCreateUser())
                    .inputTokenCount(vo.getInputTokenCount().intValue())
                    .outputTokenCount(vo.getOutputTokenCount().intValue())
                    .tokenCount((int) (vo.getInputTokenCount() + vo.getOutputTokenCount()))
                    .usageDesc(vo.getDesc())
                    .recordTime(new Date())
                    .modelName(StringUtils.isBlank(normalChatDTO.getModelName()) ? defaultModelName : normalChatDTO.getModelName())
                    .build();
            lLmUsageRecordsMapper.insert(records);
            LLMRecords llmRecord = LLMRecords.builder()
                    .outputContent(vo.getContent())
                    .updateTime(new Date())
                    .build();
            lLmRecordsMapper.update(llmRecord, Wrappers.<LLMRecords>lambdaQuery().eq(LLMRecords::getChatId, chatId));
        }
        return chatId;
    }

    @Override
    public R<LLMResponseVO> chat(NormalChatDTO normalChatDTO) {
        LLMStrategy llm = getLLM(normalChatDTO);
        normalChatDTO.setCreateDate(new Date());
        R<LLMResponseVO> chat = llm.chat(normalChatDTO);
        if (chat.getIsSuccess()) {
            LLMResponseVO data = chat.getData();
            record(normalChatDTO, data, null, false);
        }
        return chat;
    }


    @Override
    public ResponseBodyEmitter streamChat(StreamChatDTO streamChatDTO, HttpServletResponse response, Consumer<LLMStreamResultDTO> resultHandler) {
        LLMStrategy llm = getLLM(streamChatDTO);
        return llm.streamChat(streamChatDTO, response, result -> {
            if (result.isStatus()) {
                record(streamChatDTO, result.getLlmResponseVO(), streamChatDTO.getChatId(), true);
            } else {
                log.warn("流式对话失败，chatId:{},返回内容:{}", streamChatDTO.getChatId(), result);
            }
            resultHandler.accept(result);
        });
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        initStrategy();
    }

}
